import os
from shapely.geometry import CAP_STYLE, JOIN_STYLE
from shapely.ops import unary_union
import geopandas as gpd
import datetime
import networkx as nx
from utils.basic_function import *
from settings import *
from merge.utils.DataProvider import DataSetProvider


class FuseTool:
    NEW_LDBM_COL = 'match_ldbm'

    def __init__(self):
        pass

    def create_new_ldbm_col(self, rp_table):
        with conn:
            with conn.cursor() as cursor:
                sql = f"""
                            alter table {rp_table.name} add column if not exists {self.NEW_LDBM_COL} integer;
                            update {rp_table.name} set {self.NEW_LDBM_COL} = {rp_table.pk};"""
                cursor.execute(sql)

    def create_join_table(self, join_table):
        cur = conn.cursor()
        sql = f"SELECT postgis_extensions_upgrade()"
        cur.execute(sql)
        if type(join_table) != str:
            # 按照FuseRoadTable设计的Fuse表--------------------------------------------------------------
            join_table.create_table(drop=True)
            sql = f"alter table {join_table.name} add column if not exists {self.NEW_LDBM_COL} integer"
            cur.execute(sql)
            conn.commit()
        else:
            #### 原本只有三个字段的表------------------------------------------------------------------------
            # 创建融合路网表
            sql = f"""
                        drop table if exists {join_table};
                        create table if not exists {join_table} ({PK_COL} serial, {self.NEW_LDBM_COL} integer);
                        alter table {join_table} add column if not exists iteration integer;
                        """
            cur.execute(sql)
            conn.commit()
            postgis_add_geomcol(join_table,
                                conn,
                                SRID,
                                "linestring",
                                print_text=False)

    def insert_no_match_road(self, join_table, match_table, rp_table):
        cur = conn.cursor()
        # 恢复原始路段编码，加入没有匹配结果的上报路段
        sql = f"""insert into {join_table}({self.NEW_LDBM_COL}, geom) (select {self.NEW_LDBM_COL}, geom from {rp_table}
                        where {self.NEW_LDBM_COL} not in (select {MATCH_COL} from {match_table}));
                    """
        cur.execute(sql)
        conn.commit()


class ChangeDetectionTool(FuseTool):
    def __init__(self,extent):
        self.extent = extent
        pass

    def change_detection_workflow(self,
                                src_rp_name,
                                rs_layer,
                                rp_layer,
                                output_layer,
                                sourceType="DATABASE"):
        """
        变化检测法路网融合工作流

        Args:
            src_rp_name (str): 当数据源为数据库时，表示上报路网图层名
            rs_layer (database_table.RemoteSensingRoadTable): 遥感路网表结构
            rp_layer (database_table.ReportRoadTable): 上报路网表结构
            output_layer (database_table.FuseRoadTable): 合并路网表结构
            save (bool, optional): 是否保存融合过程图像. Defaults to False.
            sourceType (str, optional): 指定数据源类型，可选值：["DATABASE","SHP"]. Defaults to "DATABASE".

        Raises:
            NotImplementedError: [description]
        """
                
        print(f"正在处理{rs_layer.name}")
        # * 1.记录程序开始时间
        time1 = datetime.datetime.now()

        ## * 2. 设置运行环境

        # * 2.1 构建融合数据集
        try:
            if sourceType == "DATABASE":
                # 从数据库获取数据
                data = DataSetProvider(rs_layer,
                                       rp_layer,
                                       self.extent,
                                       THRESHOLD_MATCH_LENGTH,
                                       THRESHOLD_EXTEND_TOLERANCE,
                                       THRESHOLD_BUFFER_COVER,
                                       SRID)
            else:
                raise NotImplementedError("未知数据源类型")
        except FileNotFoundError as fileNotFoundError:
            print(str(fileNotFoundError))
            return
        except NotImplementedError as notImplementedError:
            print(str(notImplementedError))
            return
        except ConnectionError as connectionerror:
            print(str(connectionerror))
            return

        self.create_new_ldbm_col(rp_layer)
        self.create_join_table(output_layer)  # todo

        # * 3. 开始数据融合
        data.join()

        # * 4. 记录融合结束时间
        time2 = datetime.datetime.now()
        print(f"{rs_layer.name}图层匹配耗时：{(time2 - time1).total_seconds()}s")

        # * 5. 融合结果上传数据库

        # 将结果上传到数据库，此时还是分段的
        data.Post_JoinResult_To_Database(src_table_name=src_rp_name.name,
                                         target_table_name=output_layer.name,
                                         con=engine,
                                         columns=[self.NEW_LDBM_COL, 'geom'],
                                         if_exists='replace')


class BufferIterationTool(FuseTool):
    def __init__(self, buffer_dist=5, Iteration_num=1):
        super().__init__()
        self.buffer_dist = buffer_dist
        self.Iteration_num = Iteration_num

    def buffer_iteration_workflow(self, source_rp_table, rp_table, match_table, join_table):
        """
        根据一定距离构建上报路网和匹配路段的缓冲区，提取融合结果
        Args:
            source_rp_table: 原始上报路网表（未打断，属性完整，用于填充属性）
            rp_table: 上报路网表，必备字段（gid int, new_ldbm integer, ldbm text, geom linestring）
            match_table: 匹配路段表，必备字段（gid int, match_road integer, geom linestring）
            join_table: 融合路网表，必备字段（gid int, ldbm text, geom linestring）
                                    新增字段, 迭代次数（iteration_num int）
        """
        # 添加临时字段iteration
        join_table.iteration = "iteration"
        join_table.dtype_dict[join_table.iteration] = "smallint"
        join_table.match_ldbm = self.NEW_LDBM_COL
        join_table.dtype_dict[join_table.match_ldbm] = "integer"
        join_table.create_table(drop=True)

        buffer_dist = self.buffer_dist

        # 融合路网
        self.buffer_iteration(rp_table, match_table, join_table, buffer_dist)

        # 将无匹配结果道路插入到融合结果当中
        self.insert_no_match_road(join_table, match_table, rp_table)
        # 匹配程度评价，单独输出 or 在属性中标注？

        # 合并同路段编码道路
        self.road_merge_by_ldbm(join_table, rp_table)
        # 属性合并
        self.attribute_merge(source_rp_table, join_table)

    def buffer_iteration(self, rp_table, match_table, join_table, buffer_dist):
        """
        迭代融合路网，将结果保存在join_table当中
        join_table: (必要字段) ldbm text, gid int, iteration int, geom linestring
        """
        sql = f"select distinct {match_table.match_road} from {match_table.name} order by {match_table.match_road}"
        rp_pk_list = pd.read_sql(sql, conn).values.flatten()
        for rp_pk in rp_pk_list:
            print(rp_pk)
            # 获取匹配的遥感路段
            match_gdf = match_table.select_by_attr(match_table.match_road, rp_pk)
            match_geom = unary_union(match_gdf[match_table.geom])
            # 获取上报路段
            rp_gdf = rp_table.select_by_attr(self.NEW_LDBM_COL, rp_pk)
            assert rp_gdf.shape[0] == 1
            rp_geom = rp_gdf.loc[0, rp_table.geom]
            if match_geom.geometryType() == 'LineString':
                match_geom = self.match_road_alignment(rp_geom, match_geom)  # 对齐路段
            # 缓冲区融合
            try:
                fuse_geom = self.join_ldbm(rp_geom, match_geom, buffer_dist * 2)
                i = 1
                fuse_pk = join_table.insert({join_table.geom: fuse_geom, join_table.iteration: i, self.NEW_LDBM_COL: rp_pk})
                buffer_sim = self.compute_buffer_similarity(fuse_geom, match_geom, buffer_dist * 2)
                buffer_sim1 = 0
                # 迭代
                while buffer_sim < 0.7 and i < self.Iteration_num and abs(buffer_sim - buffer_sim1) > 0.1:
                    i = i + 1
                    buffer_sim1 = buffer_sim
                    fuse_geom = self.join_ldbm(fuse_geom, match_geom, buffer_dist)
                    join_table.update_attr(join_table.pk, fuse_pk,
                                           {join_table.geom: fuse_geom, join_table.iteration: i})
                    buffer_sim = self.compute_buffer_similarity(fuse_geom, match_geom, buffer_dist * 2)
            except Exception as e:  # 不成功则保留上报路网
                print(e)
                with conn:
                    join_table.delete_by_attr(self.NEW_LDBM_COL, rp_pk, auto_commit=False)
                    join_table.insert({join_table.geom: rp_geom, self.NEW_LDBM_COL: rp_pk, join_table.iteration: 0})

    def join_ldbm(self, rp_geom, match_geom, buffer_dist):
        """
        融合单个路段，返回融合结果
        """
        # 上报路段缓冲区
        rp_buffer = rp_geom.buffer(buffer_dist)
        # 匹配路段缓冲区
        match_buffer = match_geom.buffer(buffer_dist)
        # 合并上报路网和遥感路网的缓冲区
        buffer_union = Polygon(unary_union([match_buffer, rp_buffer]).exterior.simplify(0.5))
        # 缓冲区融合
        fuse_geom = self.create_road_by_buffer(buffer_union, rp_geom)
        return fuse_geom

    def create_road_by_buffer(self, buffer_union, rp_geom):
        #  提取合并缓冲区外边界及其中心线
        sql = f"""-- 中心线 需要create extension postgis_sfcgal;
            select ST_ApproximateMedialAxis({geom_to_ewkt(buffer_union, SRID)}) as geom;
            """
        gdf = geopd.read_postgis(sql, conn)
        medial_axis = gdf.loc[0, "geom"]

        st_coord = rp_geom.coords[0]
        ed_coord = rp_geom.coords[-1]
        new_geom = self.build_network_from_geom(medial_axis, st_coord, ed_coord)

        # 将融合路段端点捕捉到上报路段端点
        new_geom = rubber_snap_within_distance(Point(st_coord), new_geom, False, 50)
        new_geom = rubber_snap_within_distance(Point(ed_coord), new_geom, True, 50)
        return new_geom

    def build_network_from_geom(self, geom, source_coord, target_coord):
        if geom.geometryType() != 'MultiLineString':  # 融合结果为单条线段，调整方向后直接返回
            if cal_dist(source_coord, geom.coords[0]) > cal_dist(target_coord, geom.coords[0]):
                return reverse_linestring(geom)
            else:
                return geom

        # 构建无向图
        G = nx.Graph()
        node_arr = np.array([])
        edge_list = []

        for ls in geom:
            st_coord = ls.coords[0]
            snodeid, node_arr = self.find_node(node_arr, st_coord)
            ed_coord = ls.coords[-1]
            enodeid, node_arr = self.find_node(node_arr, ed_coord)
            edge_list.append((snodeid, enodeid))

        source_id = self.find_nearest_node(node_arr, source_coord)
        target_id = self.find_nearest_node(node_arr, target_coord)
        G.add_edges_from(edge_list)
        path_node_list = nx.dijkstra_path(G, source=source_id, target=target_id)

        new_geom = LineString(node_arr[path_node_list])
        return new_geom

    @staticmethod
    def find_nearest_node(arr, coord):
        dist = (arr[:, 0] - coord[0])**2 + (arr[:, 1] - coord[1])**2
        return dist.argmin()

    def find_node(self, node_arr, coord):
        """
        在node_arr里寻找是否已有coord，如果有则返回其下标，如果没有则在末尾插入并返回其下标
        """
        if len(node_arr) == 0:
            node_arr = np.array(coord).reshape(-1, 2)
            return 0, node_arr
        idx = self.find_nearest_node(node_arr, coord)
        if is_same_point(coord, node_arr[idx]):
            return idx, node_arr
        else:
            node_arr = np.vstack([node_arr, coord])
            return node_arr.shape[0] - 1, node_arr

    def match_road_alignment(self, rp_geom, rs_geom):
        """
        路段对齐
        如果只有一段匹配路段，长度相近，端点距离在一定距离范围
        则将匹配路段按照两端点，对齐到上报路段
        若重叠度变小
        Returns:

        """
        # 长度是否相近,小于上报路段的百分之15
        rs_len = rs_geom.length
        rp_len = rp_geom.length
        len_diff = abs(rp_len - rs_len)
        if len_diff < 0.15 * rp_len and len_diff < 100:
            # 端点距离判断
            # rp起点
            rp_st = Point(rp_geom.coords[0])
            rp_ed = Point(rp_geom.coords[-1])
            match_st = Point(rs_geom.coords[0])
            match_ed = Point(rs_geom.coords[-1])

            st_dist = min(rp_st.distance(match_st), rp_st.distance(match_ed))
            ed_dist = min(rp_ed.distance(match_st), rp_ed.distance(match_ed))

            if 10 < st_dist < 50 and 10 < ed_dist < 50:
                # 按照端点进行橡皮变换（线性）
                if rp_st.distance(match_st) < rp_st.distance(match_ed):
                    tag = 1
                else:
                    tag = 2
                new_geom = self.rubber_transform(rp_geom, rs_geom, tag)

                buffer_sim = self.compute_buffer_similarity(rp_geom, rs_geom, 30)
                buffer_sim_new = self.compute_buffer_similarity(rp_geom, new_geom, 30)
                if buffer_sim_new > buffer_sim:
                    rs_geom = new_geom
        return rs_geom

    def rubber_transform(self, rp_geom, rs_geom, tag):
        rp_x1, rp_y1 = rp_geom.coords[0]
        rp_xn, rp_yn = rp_geom.coords[-1]

        if tag != 1:
            rs_geom = reverse_linestring(rs_geom)

        rs_x1, rs_y1 = rs_geom.coords[0]
        rs_xn, rs_yn = rs_geom.coords[-1]

        dx1 = rp_x1 - rs_x1
        dy1 = rp_y1 - rs_y1
        dxn = rp_xn - rs_xn
        dyn = rp_yn - rs_yn

        rs_coords_list = list(rs_geom.coords)
        point_num = len(rs_coords_list)

        new_xy = []
        for index, pxy in enumerate(rs_coords_list):
            px = dx1 * ((point_num - index) / point_num) + pxy[0]
            py = dy1 * ((point_num - index) / point_num) + pxy[1]

            px = dxn * (index / point_num) + px
            py = dyn * (index / point_num) + py
            new_xy.append((px, py))

        return LineString(new_xy)

    @staticmethod
    def compute_buffer_similarity(rp_geom, rs_geom, buff_dist):
        """
        计算缓冲区相交部分的面积
        """
        rp_buff = rp_geom.buffer(buff_dist,
                                 cap_style=CAP_STYLE.square,
                                 join_style=JOIN_STYLE.round)
        rs_buff = rs_geom.buffer(buff_dist,
                                 cap_style=CAP_STYLE.square,
                                 join_style=JOIN_STYLE.round)
        intersect_area = rp_buff.intersection(rs_buff).area
        return intersect_area / rs_buff.area

    def insert_no_match_road(self, join_table, match_table, rp_table):
        """
        恢复原始路段编码，加入没有匹配结果的上报路段
        """
        with conn:
            with conn.cursor() as cursor:
                sql = f"""
                    insert into {join_table.name}({self.NEW_LDBM_COL}, {join_table.geom}, {join_table.length})  (
                        select {self.NEW_LDBM_COL}, {rp_table.geom}, st_length({rp_table.geom}) from {rp_table.name}
                        where {self.NEW_LDBM_COL} not in 
                            (select {match_table.match_road} from {match_table.name})
                    )
                    """
                cursor.execute(sql)

    def road_merge_by_ldbm(self, join_table, rp_table):
        with conn:
            with conn.cursor() as cursor:
                sql = f"""
                     update {join_table.name} set {join_table.segment_code} = a.{rp_table.segment_code} 
                     from {rp_table.name} a where {join_table.name}.{self.NEW_LDBM_COL} = a.{self.NEW_LDBM_COL}"""
                cursor.execute(sql)

        # 按照上报路网路段编码，合并分割的路段
        sql = f"""select {join_table.segment_code}, count(*) as {join_table.iteration} from {join_table.name} 
            group by {join_table.segment_code} having count(*) > 1;"""
        multi_df = pd.read_sql(sql, conn)
        for idx in multi_df.index:
            ldbm = multi_df.loc[idx, join_table.segment_code]
            with conn:
                sql = f"""select st_linemerge(ST_SnapToGrid(st_collect(geom), 0.1)) as geom 
                            from {join_table.name} where {join_table.segment_code}='{ldbm}'"""
                new_geom = geopd.read_postgis(sql, conn).loc[0, "geom"]
                join_table.delete_by_attr(join_table.segment_code, ldbm, auto_commit=False)
                join_table.insert({join_table.geom: new_geom, join_table.segment_code: ldbm})

    @staticmethod
    def attribute_merge(rp_table, join_table):
        rp_col = list(rp_table.dtype_dict.keys())
        for col in [rp_table.pk, rp_table.length, rp_table.geom, rp_table.segment_code]:
            rp_col.remove(col)

        # 合并上报路网属性+融合路网属性（geom、iteration）
        with conn:
            with conn.cursor() as cur:
                sql = f""" 
                    update {join_table.name} a
                    set {",".join([x + "=b." + x for x in rp_col])}
                    from {rp_table.name} b
                    where a.{join_table.segment_code} = b.{rp_table.segment_code};
                    """
                cur.execute(sql)
